from absl import app, flags
import datetime
import logging
import math
import traceback
import MySQLdb

import coin.base.database.util as db_util

FLAGS = flags.FLAGS


def cleanup_latency_by_info_id(connection, min_latency_info_id, max_latency_info_id, iter_num):
  cursor = connection.cursor()
  sql = """
DELETE FROM StrategyLatencyStatsInfo 
WHERE strategy_latency_info_id>=%s 
AND strategy_latency_info_id<=%s
""" % (min_latency_info_id, max_latency_info_id)
  cursor.execute(sql)
  sql = """
DELETE FROM StrategyLatencyStats 
WHERE strategy_latency_info_id >= %s 
AND strategy_latency_info_id <= %s
""" % (min_latency_info_id, max_latency_info_id)
  cursor.execute(sql)
  connection.commit()
  return


def cleanup_active_latency(connection, start_date, end_date):
  sql = 'DELETE FROM ActiveStrategyLatencyStats WHERE stats_timestamp < "%s"' % end_date
  assert end_date is not None
  if start_date is not None:
    sql += ' AND stats_timestamp >= "%s"' % start_date
  connection.cursor().execute(sql)
  connection.commit()
  return


def cleanup_latency_db(start_date, end_date):
  db_config = db_util.read_db_config(FLAGS.latency_db_config)
  db_connection = MySQLdb.connect(**db_config)
  cursor = db_connection.cursor()
  sql = """
SELECT strategy_latency_info_id FROM StrategyLatencyStatsInfo 
WHERE stats_timestamp <= '%s'
""" % end_date
  if start_date:
    sql += " AND stats_timestamp >= '%s'" % start_date
  sql += " ORDER BY strategy_latency_info_id"
  cursor.execute(sql)
  results = cursor.fetchall()
  results = [res[0] for res in results]
  rows_to_read = 10000
  iter_cnt = math.ceil(len(results) / float(rows_to_read))
  logging.info("Start cleanup from %s to %s" % (start_date, end_date))
  for iter_num in range(iter_cnt):
    min_latency_info_id = results[rows_to_read * iter_num]
    max_latency_info_id = results[rows_to_read * (iter_num + 1) if iter_num != iter_cnt - 1 else -1]
    try:
      cleanup_latency_by_info_id(db_connection, min_latency_info_id, max_latency_info_id, iter_num)
      logging.info("Latency DB cleanup done(%s/%s)" % (iter_num, iter_cnt - 1))
    except Exception as e:
      db_connection.rollback()
      logging.error(traceback.format_exc())
  cleanup_active_latency(db_connection, start_date, end_date)
  db_connection.close()


def main(_):
  if FLAGS.end_date is None:
    end_date = datetime.datetime.today().date() - datetime.timedelta(days=3)
  else:
    end_date = datetime.datetime.strptime(FLAGS.end_date, '%Y%m%d').date()
  start_date = FLAGS.start_date
  if start_date is not None:
    start_date = datetime.datetime.strptime(start_date, '%Y%m%d').date()
    assert start_date <= end_date, (start_date, end_date)
  cleanup_latency_db(start_date, end_date)


if __name__ == '__main__':
  flags.DEFINE_string('latency_db_config',
                      '../../coin_deploy/support_pta/config/pta_db_config.json',
                      'path of latency db config.')

  flags.DEFINE_string('start_date', None, 'yyyymmdd')
  flags.DEFINE_string('end_date', None, 'yyyymmdd')
  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
